import os
os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'

import torch
import torch.nn as nn
import torch.optim as optim
import h5py
import numpy as np
from torch.utils.data import Dataset, DataLoader
import os
from tqdm import tqdm
import matplotlib.pyplot as plt
import pandas as pd

class VoltageProcessor:
    @staticmethod
    def calculate_calibration_matrix(Vcav_I, Vcav_Q, Vfor_I, Vfor_Q, Vrefl_I, Vrefl_Q):
        """计算校准矩阵和校准后的电压值"""
        # 构建校准矩阵
        A_matrix = np.column_stack((
            np.append(Vfor_I, Vfor_Q),
            np.append(-Vfor_Q, Vfor_I),
            np.append(Vrefl_I, Vrefl_Q),
            np.append(-Vrefl_Q, Vrefl_I)
        ))
        y_array = np.append(Vcav_I, Vcav_Q)

        try:
            C_array = np.matmul(np.linalg.pinv(A_matrix, rcond=1e-5), y_array.transpose())
            return C_array
        except np.linalg.LinAlgError as e:
            print(f"SVD计算失败: {e}")
            return None

class WaveformDataset(Dataset):
    def __init__(self, file_path, sample_ratio=0.2):
        """初始化数据集
        
        Args:
            file_path (str): HDF5文件路径
            sample_ratio (float): 采样比例
        """
        self.file_path = file_path
        
        # 读取数据
        with h5py.File(file_path, 'r') as f:
            # 获取数据集数量
            self.n_sets = len([key for key in f.keys() if key.startswith('set_')])
            n_samples = int(self.n_sets * sample_ratio)
            
            # 随机选择数据集
            np.random.seed(42)
            selected_indices = np.random.choice(self.n_sets, n_samples, replace=False)
            selected_indices.sort()
            
            # 初始化数据列表
            self.features = []
            self.labels = []
            
            for idx in tqdm(selected_indices, desc="加载数据"):
                set_name = f'set_{idx}'
                dataset = f[set_name]
                
                # 读取波形数据
                Vc = dataset['Vc_real'][()] + 1j * dataset['Vc_imag'][()]
                Vf_star = dataset['Vf_star_real'][()] + 1j * dataset['Vf_star_imag'][()]
                Vr_star = dataset['Vr_star_real'][()] + 1j * dataset['Vr_star_imag'][()]
                
                # 读取标量参数
                a = float(dataset['a_real'][()]) + 1j * float(dataset['a_imag'][()])
                b = float(dataset['b_real'][()]) + 1j * float(dataset['b_imag'][()])
                c = float(dataset['c_real'][()]) + 1j * float(dataset['c_imag'][()])
                d = float(dataset['d_real'][()]) + 1j * float(dataset['d_imag'][()])
                
                # 构建特征向量
                features = np.concatenate([
                    np.abs(Vc), np.angle(Vc),
                    np.abs(Vf_star), np.angle(Vf_star),
                    np.abs(Vr_star), np.angle(Vr_star)
                ])
                
                # 构建标签向量（只包含a,b,c,d）
                labels = np.array([
                    a.real, a.imag,
                    b.real, b.imag,
                    c.real, c.imag,
                    d.real, d.imag
                ])
                
                self.features.append(features)
                self.labels.append(labels)
            
            # 转换为numpy数组
            self.features = np.array(self.features)
            self.labels = np.array(self.labels)
            
            # 转换为PyTorch张量
            self.features = torch.FloatTensor(self.features)
            self.labels = torch.FloatTensor(self.labels)
    
    def __len__(self):
        return len(self.features)
    
    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]

class ABCDNeuralNetwork(nn.Module):
    def __init__(self, input_size):
        super(ABCDNeuralNetwork, self).__init__()
        
        # 五层神经网络
        self.layer1 = nn.Linear(input_size, 512)
        self.layer2 = nn.Linear(512, 256)
        self.layer3 = nn.Linear(256, 128)
        self.layer4 = nn.Linear(128, 64)
        self.layer5 = nn.Linear(64, 8)  # 只输出8个值：a,b,c,d的实部和虚部
        
        # 激活函数
        self.relu = nn.ReLU()
        
    def forward(self, x):
        x = self.relu(self.layer1(x))
        x = self.relu(self.layer2(x))
        x = self.relu(self.layer3(x))
        x = self.relu(self.layer4(x))
        x = self.layer5(x)
        return x

def calculate_errors(output, Vc, Vf, Vr, dt, omega_05):
    """计算四个误差函数"""
    # 将输出转换为复数形式
    a = output[:, 0] + 1j * output[:, 1]
    b = output[:, 2] + 1j * output[:, 3]
    c = output[:, 4] + 1j * output[:, 5]
    d = output[:, 6] + 1j * output[:, 7]
    
    # 计算X和Y
    Vcav_I = Vc.real
    Vcav_Q = Vc.imag
    Vfor_I = Vf.real
    Vfor_Q = Vf.imag
    Vrefl_I = Vr.real
    Vrefl_Q = Vr.imag
    
    # 对每个批次样本计算X和Y
    X_list = []
    Y_list = []
    for i in range(len(Vcav_I)):
        C_array = VoltageProcessor.calculate_calibration_matrix(
            Vcav_I[i], Vcav_Q[i],
            Vfor_I[i], Vfor_Q[i],
            Vrefl_I[i], Vrefl_Q[i]
        )
        if C_array is not None:
            X = complex(C_array[0], C_array[1])
            Y = complex(C_array[2], C_array[3])
            X_list.append(X)
            Y_list.append(Y)
        else:
            X_list.append(0+0j)
            Y_list.append(0+0j)
    
    X = torch.tensor(X_list, device=output.device)
    Y = torch.tensor(Y_list, device=output.device)
    
    # 计算f_probe
    f_probe = Vc - Vf - Vr
    
    # 计算f_c
    dVc_dt = torch.gradient(torch.abs(Vc), dim=1)[0]
    theta = torch.angle(Vc)
    phi = torch.angle(Vf)
    f_c = dVc_dt + torch.abs(Vc) * omega_05 - omega_05 * 2 * torch.abs(Vf) * torch.cos(theta - phi)
    
    # 计算f_X和f_Y
    f_X = X - (a + c)
    f_Y = Y - (b + d)
    
    return f_probe, f_c, f_X, f_Y

def train_model(model, train_loader, optimizer, device, num_epochs=100):
    """训练模型"""
    model.train()
    for epoch in range(num_epochs):
        total_loss = 0
        for batch_features, batch_labels in train_loader:
            batch_features = batch_features.to(device)
            batch_labels = batch_labels.to(device)
            
            # 前向传播
            outputs = model(batch_features)
            
            # 计算误差
            f_probe, f_c, f_X, f_Y = calculate_errors(
                outputs,
                batch_features[:, :1000],  # Vc
                batch_features[:, 1000:2000],  # Vf
                batch_features[:, 2000:3000],  # Vr
                0.001,  # dt
                2 * np.pi * 50  # omega_05
            )
            
            # 计算总损失
            loss = (torch.mean(torch.abs(f_probe)) + 
                   torch.mean(torch.abs(f_c)) +
                   torch.mean(torch.abs(f_X)) + 
                   torch.mean(torch.abs(f_Y)))
            
            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        # 打印训练进度
        if (epoch + 1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(train_loader):.4f}')

def save_predictions(model, test_loader, device, save_path):
    """保存预测结果到CSV文件
    
    Args:
        model: 训练好的模型
        test_loader: 测试数据加载器
        device: 计算设备
        save_path: 保存路径
    """
    model.eval()  # 设置为评估模式
    
    # 存储预测值和实际值
    predictions = []
    actuals = []
    
    with torch.no_grad():
        for features, labels in test_loader:
            features = features.to(device)
            labels = labels.to(device)
            
            # 获取预测值
            outputs = model(features)
            
            # 将预测值和实际值添加到列表中
            predictions.append(outputs.cpu().numpy())
            actuals.append(labels.cpu().numpy())
    
    # 合并所有批次的结果
    predictions = np.vstack(predictions)
    actuals = np.vstack(actuals)
    
    # 创建DataFrame
    results_df = pd.DataFrame({
        'a_true_real': actuals[:, 0],
        'a_true_imag': actuals[:, 1],
        'a_pred_real': predictions[:, 0],
        'a_pred_imag': predictions[:, 1],
        'b_true_real': actuals[:, 2],
        'b_true_imag': actuals[:, 3],
        'b_pred_real': predictions[:, 2],
        'b_pred_imag': predictions[:, 3],
        'c_true_real': actuals[:, 4],
        'c_true_imag': actuals[:, 5],
        'c_pred_real': predictions[:, 4],
        'c_pred_imag': predictions[:, 5],
        'd_true_real': actuals[:, 6],
        'd_true_imag': actuals[:, 7],
        'd_pred_real': predictions[:, 6],
        'd_pred_imag': predictions[:, 7]
    })
    
    # 计算相对误差
    for param in ['a', 'b', 'c', 'd']:
        true_complex = results_df[f'{param}_true_real'] + 1j * results_df[f'{param}_true_imag']
        pred_complex = results_df[f'{param}_pred_real'] + 1j * results_df[f'{param}_pred_imag']
        
        # 计算相对误差
        relative_error = np.abs(pred_complex - true_complex) / np.abs(true_complex) * 100
        results_df[f'{param}_relative_error_%'] = relative_error
    
    # 保存到CSV
    results_df.to_csv(save_path, index=False)
    print(f"\n预测结果已保存到: {save_path}")
    
    # 打印统计信息
    print("\n预测误差统计:")
    for param in ['a', 'b', 'c', 'd']:
        errors = results_df[f'{param}_relative_error_%']
        print(f"\n{param}参数:")
        print(f"平均相对误差: {errors.mean():.2f}%")
        print(f"最大相对误差: {errors.max():.2f}%")
        print(f"中位相对误差: {errors.median():.2f}%")

def main():
    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    
    # 加载数据
    data_path = r'E:\MLdata\SimData\coupling_simulation_data.h5'
    dataset = WaveformDataset(data_path)
    
    # 分割训练集和测试集
    train_size = int(0.8 * len(dataset))
    test_size = len(dataset) - train_size
    train_dataset, test_dataset = torch.utils.data.random_split(
        dataset, [train_size, test_size]
    )
    
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
    
    # 创建模型
    input_size = dataset.features.shape[1]
    model = ABCDNeuralNetwork(input_size).to(device)
    
    # 定义优化器
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    # 训练模型
    print("开始训练...")
    train_model(model, train_loader, optimizer, device)
    
    # 保存模型
    save_dir = os.path.dirname(data_path)
    model_path = os.path.join(save_dir, 'abcd_model.pth')
    torch.save(model.state_dict(), model_path)
    print(f"模型已保存到: {model_path}")
    
    # 保存预测结果
    results_path = os.path.join(save_dir, 'abcd_predictions.csv')
    save_predictions(model, test_loader, device, results_path)

if __name__ == '__main__':
    main() 